#!/usr/bin/python
from transformers import AutoTokenizer, BartForConditionalGeneration
tokenizer = AutoTokenizer.from_pretrained("IDEA-CCNL/Randeng-BART-139M-QG-Chinese",additional_special_tokens=["<ans>"])
model = BartForConditionalGeneration.from_pretrained("IDEA-CCNL/Randeng-BART-139M-QG-Chinese")

context = "知识：1939年9月1日德国入侵波兰后，第二次世界大战开始，华沙一直被保卫到9月27日。波兰中部，包括华沙，都在德国纳粹殖民地政府总政府的统治下。所有的高等教育机构都立即关闭，华沙的犹太人口——几十万，约占城市的 <ans> ——全部涌入华沙的贫民区。回答：30%"
inputs = tokenizer.encode_plus(
            context,
            max_length=448,
            padding="max_length",
            truncation=True,
            return_tensors='pt'
        )
output = model.generate(                
        input_ids=inputs['input_ids'],
        attention_mask=inputs['attention_mask'],
        do_sample=True,
        num_beams=5,
        max_length=64,
        top_p = 0.9,
    )
result = tokenizer.batch_decode(output,clean_up_tokenization_spaces=True, skip_special_tokens=True)[0].replace("问题:","")
print(result)
# 问题:华沙的犹太人口占城市的百分之多少?
